# 开发者：郭同源
# 开发时间：2021/10/31 10:40
import os

from my_model import MyModel
from dataLoader import MyData

import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torch import nn,optim

DEVICE = 'cpu'


class Train():
    def __init__(self,root,weight):
        self.writer = SummaryWriter('logs')
        self.train_dataset = MyData(root)
        self.test_dataset = MyData('data/test/test')
        self.train_dataloader = DataLoader(self.train_dataset,batch_size=50,shuffle=True)
        self.test_dataloader = DataLoader(self.test_dataset,batch_size=50,shuffle=True)


        self.model = MyModel().to(DEVICE)
        if os.path.exists(weight):
            self.model.load_state_dict(torch.load(weight))
        self.opt = optim.Adam(self.model.parameters())

        self.label_loss = nn.BCEWithLogitsLoss()
        self.position_loss = nn.MSELoss()
        self.sort_loss = nn.CrossEntropyLoss()


        self.train = True
        self.test = True


    def __call__(self):
        index1,index2 = 0,0
        for epoch in range(20):
            for i, (img, label, position, sort) in enumerate(self.train_dataloader):
                self.model.train()
                img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                # print(label.shape)
                out_label, out_position, out_sort = self.model(img)

                # print(out_label.shape)
                position_loss = self.position_loss(out_position, position)
                out_sort = out_sort[torch.where(sort >= 0)]
                sort = sort[torch.where(sort >= 0)]
                sort_loss = self.sort_loss(out_sort, sort)
                label_loss = self.label_loss(out_label, label)

                train_loss = 0.2 * label_loss + position_loss * 0.6 + 0.2 * sort_loss
                # print(train_loss)

                self.opt.zero_grad()
                train_loss.backward()
                self.opt.step()

                if i % 10 == 0:
                    print(f'train_loss {i}===>>', train_loss.item())
                    self.writer.add_scalar('train_loss', train_loss, index1)
                    index1 += 1
            # torch.save(self.model.state_dict(),'param/model{}.pth'.format(epoch))
            torch.save(self.model.state_dict(),'model/model{}.pth'.format(epoch))

            self.model.eval()
            sum_acc = 0
            for i, (img, label, position, sort) in enumerate(self.test_dataloader):
                img, label, position, sort = img.to(DEVICE), label.to(DEVICE), position.to(DEVICE), sort.to(DEVICE)
                out_label, out_position, out_sort = self.model(img)


                position_loss = self.position_loss(out_position, position)
                out_sort = out_sort[torch.where(sort >= 0)]
                sort = sort[torch.where(sort >= 0)]
                sort_loss = self.sort_loss(out_sort, sort)
                label_loss = self.label_loss(out_label, label)

                test_loss = 0.2 * label_loss + position_loss * 0.6 + 0.2 * sort_loss

                sort_acc = torch.mean(torch.eq(sort, torch.argmax(torch.softmax(out_sort, dim=1), dim=1)).float())
                sum_acc += sort_acc

                if i % 10 == 0:
                    print(f'test_loss {i}===>>', test_loss.item())
                    self.writer.add_scalar('test_loss', test_loss, index2)
                    index2 += 1

            avg_acc = sum_acc / i
            print(f'sort_acc {i}==>>', avg_acc)
            self.writer.add_scalar('avg_acc', avg_acc, epoch)













if __name__ == '__main__':
    train = Train('data/train/train',weight='model/mode4.pth')
    train()